#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

"""
    二项分布各个随机变量值概率计算, 和频率模拟.
    平均发生次数 和 期望比较.

    也可以使用scipy.stats.binom(10, 0.5) 进行二项分布计算
"""

import random
import scipy.misc
from collections import defaultdict
import matplotlib.pyplot as plt


class BinaryDistribution:

    def __init__(self, p, n, experiment_n=10**5):
        ''' n次投射作为1次试验, 发生概率为p '''
        assert p <= 1
        self.n = n
        self.p = p
        self.p_ = 1 - self.p
        # self.times_record = defaultdict(int)       #  { 击中个数: 多次试验统计个数}
        self.times_record = [0] * (self.n + 1)
        self.times_frequency = [0] * (self.n + 1)
        self.prob = [0] * (self.n + 1)
        self.experiment_n = experiment_n

        self.experiment(self.experiment_n)
        self.show_results()

    def experiment(self, experi_n):
        ''' 试验experi_n 次'''
        for i in range(experi_n):
            self.one_experiment()

        self.calc_frequency_prob()

    def one_experiment(self):
        ''' 注意一次试验是n次toss'''
        ntime = 0
        for i in range(self.n):
            result = self.one_toss()
            if result:
                ntime += 1

        self.times_record[ntime] += 1

    def calc_frequency_prob(self):
        occur = 0
        for idx, times in enumerate(self.times_record):
            self.times_frequency[idx] = times / self.experiment_n

            non_idx = self.n - idx
            self.prob[idx] = scipy.misc.comb(self.n, idx) * (self.p ** idx) * (self.p_ ** non_idx)
            occur += idx * times

        self.mean_occur     = occur / self.experiment_n      # 试验experiment_n次, 平均发生次数
        self.expected_value = self.n * self.p                     # 二项分布期望

    def one_toss(self):
        r = random.random()
        if r <= self.p:
            return 1
        else:
            return 0

    def statistics(self):
        pass

    def show_results(self):
        # for t in range(self.n):
        # print(self.times_record)
        print("frequency   : %s" % self.times_frequency)
        print("probability : %s" % self.prob)

        print('mean occur in one experiment : %s' % self.mean_occur)
        print('expected value               : %s ' % self.expected_value)


        plt.plot(range(self.n+1), self.times_frequency, 'g', label='experiment frequency')
        plt.plot(range(self.n+1), self.prob, 'r', label='real prob')
        plt.legend()        # 显示"图例"
        plt.show()


def main():
    print("start main")

    b = BinaryDistribution(0.3, 50, experiment_n=10**5)

if __name__ == "__main__":
    main()
